import sys
import os
sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
from impl.db.DAO import BaiduCrawlerDAO
from contextlib import closing
import csv
from tqdm import tqdm


def db_restore_records_from_csv(db_config: dict, csv_file: str):
    buffer_size = 100
    buffered_identities = []
    buffered_file_names = []
    buffered_urls = []
    buffered_source_urls = []
    buffered_source_image_urls = []
    buffered_captions = []
    dao = BaiduCrawlerDAO(db_config)
    with dao:
        try:
            with closing(dao.get_cursor()) as cursor, open(csv_file, 'r', newline='', encoding='utf-8') as fid:
                csv_reader = csv.reader(fid, delimiter=',')
                for row in tqdm(csv_reader):
                    if len(row) == 0:
                        continue
                    assert len(row) == 6
                    identity, file_name, url, source_url, source_image_url, caption = row
                    buffered_identities.append(identity)
                    buffered_file_names.append(file_name)
                    buffered_urls.append(url)
                    buffered_source_urls.append(source_url)
                    buffered_source_image_urls.append(source_image_url)
                    buffered_captions.append(caption)
                    if len(buffered_identities) >= buffer_size:
                        dao.insert_multiple(cursor, buffered_identities, buffered_file_names, buffered_urls, buffered_source_urls, buffered_source_image_urls, buffered_captions)
                        buffered_identities.clear()
                        buffered_file_names.clear()
                        buffered_urls.clear()
                        buffered_source_urls.clear()
                        buffered_source_image_urls.clear()
                        buffered_captions.clear()
                if len(buffered_identities) > 0:
                    dao.insert_multiple(cursor, buffered_identities, buffered_file_names, buffered_urls, buffered_source_urls, buffered_source_image_urls, buffered_captions)
            dao.commit()
        except Exception:
            dao.rollback()
            raise


import argparse


if __name__ == '__main__':
    parser = argparse.ArgumentParser('Restore records from CSV file to mysql server')
    parser.add_argument('csv_file', type=str, help='CSV file path')
    args = parser.parse_args()

    from db_utils._get_db_config import get_db_config
    db_restore_records_from_csv(get_db_config(), args.csv_file)
